# -*- coding: utf-8 -*-

# ------------------------------
# @Time    : 2023/6/15
# @Author  : gao
# @File    : tgw_api.py 
# @Project : AmazingQuant 
# ------------------------------
import tgw

from AmazingQuant.utils.date_transfer import date_to_datetime, datetime_to_int
from AmazingQuant.data_center.tgw_source.tgw_login import tgw_login


class TgwApiData(object):
    def __init__(self, ):
        self.calendar = []
        self.code_sh_list, self.code_sz_list = [], []
        self.stock_list = []
        self.code_list_hist = []

    def get_calendar(self, data_type=None):
        today = datetime_to_int()
        task_id = tgw.GetTaskID()
        tgw.SetThirdInfoParam(task_id, "function_id", "A010060001")
        tgw.SetThirdInfoParam(task_id, "start_date", "19900101")
        tgw.SetThirdInfoParam(task_id, "end_date", str(today))
        tgw.SetThirdInfoParam(task_id, "market", 'SSE')

        trade_days_df, _ = tgw.QueryThirdInfo(task_id)
        trade_days_df['TRADE_DAYS'] = trade_days_df['TRADE_DAYS'].astype(int)
        self.calendar = list(trade_days_df['TRADE_DAYS'].sort_values(ascending=True))
        if data_type == 'datetime':
            self.calendar = [date_to_datetime(str(i)) for i in self.calendar]
        return self.calendar

    def get_code_list(self, add_market=False, all_code=False, index=False):
        for market in [tgw.MarketType.kSZSE, tgw.MarketType.kSSE]:
            item = tgw.SubCodeTableItem()
            item.market = market
            item.security_code = ""
            code_table_df, error = tgw.QuerySecuritiesInfo(item)
            # print('code_table_df', code_table_df['security_type'].unique())
            if index:
                code_table_df = code_table_df[code_table_df['security_type'].isin(['01000'])]
            else:
                code_table_df = code_table_df[code_table_df['security_type'].isin(['02001', '02003', '02004', '02999'])]
            code_list = list(code_table_df['security_code'])
            if market == tgw.MarketType.kSZSE:
                self.code_sz_list = code_list
                if add_market:
                    self.code_sz_list = [i + '.SZ' for i in self.code_sz_list]
            elif market == tgw.MarketType.kSSE:
                self.code_sh_list = code_list
                if add_market:
                    self.code_sh_list = [i + '.SH' for i in self.code_sh_list]

        if all_code:
            self.stock_list = self.code_sh_list + self.code_sz_list
            return self.stock_list
        else:
            return self.code_sh_list, self.code_sz_list

    def get_code_list_hist(self):
        task_id = tgw.GetTaskID()
        tgw.SetThirdInfoParam(task_id, "function_id", "A010010006")
        tgw.SetThirdInfoParam(task_id, "start_date", "19000101")
        tgw.SetThirdInfoParam(task_id, "end_date", "20991231")
        df, error = tgw.QueryThirdInfo(task_id)
        self.code_list_hist = list(df['MARKET_CODE'])
        return self.code_list_hist


if __name__ == '__main__':
    tgw_login()
    tgw_api_object = TgwApiData()
    calendar = tgw_api_object.get_calendar()
    code_list = tgw_api_object.get_code_list()
    # item = tgw.SubCodeTableItem()
    # item.market = tgw.MarketType.kNone
    # item.security_code = ""
    # code_table_df, error = tgw.QuerySecuritiesInfo(item)
    # code_table_df['security_status'].unique()

    from AmazingQuant.environment import Environment

    # def OnResponse(result, status):
    #     # print(result)
    #     Environment.g_list_re_len -= 1
    #     Environment.data.append(result)
    #     print('g_list_re_len:', Environment.g_list_re_len)
    #     if Environment.g_list_re_len == 0:
    #         if Environment.wait_event:
    #             Environment.g_list_re_len = None
    #             Environment.wait_event.set()

    def QuerySnapshotlist(code_list):
        req_list = []
        i = 0
        for code in code_list[0][:10]:
            i += 1
            print(i)
            req_tick1 = tgw.ReqDefault()
            req_tick1.market_type = tgw.MarketType.kSSE
            req_tick1.date = 20240517
            req_tick1.begin_time = 90000000
            req_tick1.end_time = 170000000
            req_tick1.security_code = code
            req_list.append(req_tick1)
        err = None
        Environment.g_list_re_len = len(req_list)
        for req_tick in req_list:
            result, error_code = tgw.QuerySnapshot(req_tick, Environment.OnResponse)
            if error_code != tgw.ErrorCode.kSuccess:
                err = error_code
        Environment.wait_event.wait()
        import copy
        result = copy.deepcopy(Environment.data)
        Environment.refresh()
        return result, err

    index_df1, index_error1 = QuerySnapshotlist(code_list)

